#######################
### Interrupted regression results for Study 1
### Andrews, Delton, & Kline

library(MASS)
library(ggplot2)
library(sandwich)
library(lmtest)


#######################
### Auxiliary functions

## Function to recode independent variable for interrupted regression (Simonsohn 2017, p.17)
#set cutoff
cutoff_set <- .75

xvalCutoff <- function(x, cutoff = cutoff_set){
  x_lo <- ifelse(x <= cutoff, x - cutoff, 0)
  x_hi <- ifelse(x >= cutoff, x - cutoff, 0)
  d_hi <- ifelse(x >= cutoff, 1, 0)
  cbind(x_lo, x_hi, d_hi)
}

## Function to simulate expected values for interrupted regression (uses HC3 robust SEs by default)
simInterrupted <- function(model, xvals = seq(0,1,length.out = 21), xc = cutoff_set){
  ## point estimates
  beta_hat <- coef(model)
  
  ## simulate coefficients
  beta_sim <- mvrnorm(1000, mu = beta_hat, Sigma = vcovHC(model, type="HC3"))
  
  ## convert original threshold values for interrupted regression
  thresh <- cbind(1, xvalCutoff(xvals, cutoff = xc))
  xc_met <- xvals >= xc
  
  ## add cutoff observation w/ d_hi = 0
  thresh <- rbind(thresh, c(1,0,0,0))
  xvals <- c(xvals, xc)
  xc_met <- c(xc_met, FALSE)
  
  ## compute expected values
  y_hat <- as.vector(beta_hat %*% t(thresh))
  y_sim <- beta_sim %*% t(thresh)
  y_cilo <- as.vector(apply(y_sim, 2, quantile, .025))
  y_cihi <- as.vector(apply(y_sim, 2, quantile, .975))
  
  ## t-values/p-values (only used to print summary in plot)
  beta_se <- sqrt(diag(vcovHC(model, type="HC3")))
  beta_t <- beta_hat/beta_se
  beta_p <- pt(abs(beta_t), df = model$df.residual, lower.tail = FALSE) * 2
  t_val <- paste("t =", round(beta_t, 2))
  p_val <- paste("p =", round(beta_p, 3))
  p_val[p_val=="p = 0"] <- "p < 0.001"
  lab <- paste0(t_val, "\n" ,p_val)[2:3]
  
  ## create output
  list(data = data.frame(xvals, xc_met, y_hat, y_cilo, y_cihi),
       text = data.frame(xvals = c(.1,.9), y_hat = .9 * max(y_cihi), xc_met = c(F, T), lab))
}


###########
### Study 1

study1 <- read.csv("interrupted_reg_data.csv")

## check jittered scatterplot w/ loess smooth
ggplot(study1, aes(x=threshold0to1, y=gamble)) +
  geom_point() + geom_jitter() + geom_smooth() + theme_bw()

## Recode threshold variable
thresh <- xvalCutoff(study1$threshold0to1)

## Estimate OLS
m1 <- lm(study1$gamble ~ thresh)

## Check robust SEs
coeftest(m1, vcov = vcovHC(m1, type="HC3")) 

## simulate quantities of interest 
res1 <- simInterrupted(m1)

## plot results
ggplot(res1$data, aes(x = xvals, y = y_hat, group = xc_met)) + ggtitle("Study 1") +
  geom_ribbon(aes(ymin = y_cilo, ymax = y_cihi), alpha = .1) + geom_line() + 
  geom_text(aes(label = lab), data = res1$text) +
  theme_bw() + theme(legend.position="none") + labs(y = "P(Risky Contribution)", x = "Threshold") 

